import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D   # IDE 会提示未用，但必须 import
import sys, pickle

# -------------- 数据加载 -----------------
if len(sys.argv) <= 1:
    print("usage: python plot.py <dcqcn | hp95 | ...>")
    sys.exit(-1)
cc = sys.argv[1]

with open(f'data/best_alpha_2dim_{cc}_matrix.pkl', 'rb') as f:
    compare = pickle.load(f)

x = np.asarray(compare["fct_avg"]["x"])      # (Nx,)
y = np.asarray(compare["fct_avg"]["y"])      # (Ny,)
X, Y = np.meshgrid(x, y)                     # (Ny, Nx)
Z = np.asarray(compare["fct_avg"]["z"])      # (Ny, Nx)

# ----------- 找到全局极大值 ----------------
# 如果 Z 里有 NaN/inf 按需用 np.nanargmax
max_idx  = np.unravel_index(np.nanargmax(Z), Z.shape)
max_x, max_y, max_z = X[max_idx], Y[max_idx], Z[max_idx]
print(f"max@({max_x:.3f}, {max_y:.3f})  = {max_z:.3f} Gbps")

# =============== 3D 曲面图 =================
fig = plt.figure(figsize=(7, 5))
ax  = fig.add_subplot(111, projection='3d')

surf = ax.plot_surface(X, Y, Z,
                       cmap='jet',
                       edgecolor='none',
                       antialiased=True)

# 用黑边白心的小圆点 + 文字标最大值
ax.scatter(max_x, max_y, max_z,
           s=60, c='white', edgecolor='black', zorder=10)
ax.text(max_x, max_y, max_z,
        f'\n({max_x:.2f}, {max_y:.0f}, {max_z:.3f})',
        color='black', ha='center', va='bottom')

# 也可以在底面投影一条十字帮助定位
ax.plot([max_x, max_x], [max_y, max_y], [Z.min(), max_z],
        '--', color='black', lw=1)

ax.set_xlabel('Alpha')
ax.set_ylabel('Workload')
ax.set_zlabel('FCT Promotion')

# 视角：略抬高 + 越过 x 轴看，使峰值不被挡
ax.view_init(elev=32, azim=-100)

# =============== 2D 俯视热图 =================
fig2, ax2 = plt.subplots(figsize=(6, 5))

# 注意 extent 保证 x/y 数值轴对齐；origin='lower' 与 meshgrid 匹配
im = ax2.imshow(Z,
                cmap='jet',
                origin='lower',
                extent=[x.min(), x.max(), y.min(), y.max()],
                aspect='auto')

# 把最大值画出来
ax2.scatter(max_x, max_y, s=60,
            c='white', edgecolors='black')
ax2.text(max_x, max_y,
         f'{max_z:.3f}',
         color='black', ha='center', va='bottom')

ax2.set_xlabel('Alpha')
ax2.set_ylabel('Workload')
cb = fig2.colorbar(im)
cb.set_label('FCT Reduciton')

plt.tight_layout()
fig.savefig('img/best_alpha_3d.png', dpi=300, bbox_inches='tight')
fig2.savefig('img/best_alpha_topview.png', dpi=300, bbox_inches='tight')
plt.show()
